package org.jeecg.common.util.deepseek;

import java.util.HashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class SqlQueryChecker {

    public static void main(String[] args) {
        String sql = "SELECT w.*, s.SCHOOLNAME, t.TEAMMXNAME FROM web_wyjbinfo w LEFT JOIN web_hotshoolmx s ON w.ADMISSSCHOOL = s.id LEFT JOIN web_teammx t ON w.TEACHERNAME = t.id ORDER BY w.NAMETIME DESC LIMIT 2000";
        String[] allowedTables = {"web_wyjbinfo", "web_hotshoolmx", "web_teammx"};

        try {
            String validatedSql = validateAndModifySql(sql, allowedTables);
            System.out.println("Validated SQL: " + validatedSql);
        } catch (IllegalArgumentException e) {
            System.err.println("Error: " + e.getMessage());
        }
    }

    public static String validateAndModifySql(String sql, String[] allowedTables) {
        // 检查是否为查询类型
        if (!sql.trim().toUpperCase().startsWith("SELECT")) {
            throw new IllegalArgumentException("Only SELECT queries are allowed.");
        }

        // 检查表名
        for (String table : extractTableNames(sql)) {
            boolean isAllowed = false;
            for (String allowedTable : allowedTables) {
                if (table.equalsIgnoreCase(allowedTable)) {
                    isAllowed = true;
                    break;
                }
            }
            if (!isAllowed) {
                throw new IllegalArgumentException("Query contains unauthorized table: " + table);
            }
        }

        // 检查是否包含 COUNT()
        if (!sql.toUpperCase().contains("COUNT(")) {
            // 检查并添加 LIMIT 子句
            if (!sql.toUpperCase().contains("LIMIT")) {
                sql = sql.trim();
                if (sql.endsWith(";")) {
                    sql = sql.substring(0, sql.length() - 1);
                }
                sql += " LIMIT 10"; // 默认添加 LIMIT 10
            } else {
                // 确保限制在 20 条以内
                Pattern limitPattern = Pattern.compile("LIMIT\\s+(\\d+)(?:\\s*,\\s*(\\d+))?", Pattern.CASE_INSENSITIVE);
                Matcher matcher = limitPattern.matcher(sql);
                if (matcher.find()) {
                    int limitValue;
                    if (matcher.group(2) != null) {
                        // 处理 LIMIT offset, count 的情况
                        limitValue = Integer.parseInt(matcher.group(2));
                    } else {
                        // 处理 LIMIT count 的情况
                        limitValue = Integer.parseInt(matcher.group(1));
                    }

                    if (limitValue > 10) {
                        // 如果 LIMIT 值超过 20，将其改为 20
                        String replacement = "LIMIT " + (matcher.group(2) != null ? matcher.group(1) + ", 10" : "10");
                        sql = matcher.replaceFirst(replacement);
                    }
                }
            }
        }

        return sql;
    }

    private static String[] extractTableNames(String sql) {
        // 正则表达式提取表名，支持多个 JOIN 和 FROM 子句
        Pattern tablePattern = Pattern.compile(
            "(?:FROM|JOIN)\\s+([a-zA-Z0-9_]+)", // 匹配 FROM 或 JOIN 后的表名
            Pattern.CASE_INSENSITIVE
        );
        Matcher matcher = tablePattern.matcher(sql);
        Set<String> tables = new HashSet<>(); // 使用 Set 避免重复表名

        while (matcher.find()) {
            if (matcher.group(1) != null) {
                tables.add(matcher.group(1).toLowerCase()); // 统一转换为小写，避免大小写问题
            }
        }

        return tables.toArray(new String[0]); // 返回去重后的表名数组
    }
}
